#include "state_estimation/ahrs.h"
#include "kalman_filter/kalman.h"

static float ahrs_state[7] = {0.0f};
static float ahrs_P[49] = {
	1.0f,	0.0f,	0.0f,	0.0f,	0.0f,	0.0f,	0.0f,
	0.0f,	1.0f,	0.0f,	0.0f,	0.0f,	0.0f,	0.0f,
	0.0f,	0.0f,	1.0f,	0.0f,	0.0f,	0.0f,	0.0f,
	0.0f,	0.0f,	0.0f,	1.0f,	0.0f,	0.0f,	0.0f,
	0.0f,	0.0f,	0.0f,	0.0f,	0.1f,	0.0f,	0.0f,
	0.0f,	0.0f,	0.0f,	0.0f,	0.0f,	0.1f,	0.0f,
	0.0f,	0.0f,	0.0f,	0.0f,	0.0f,	0.0f,	0.1f
};

static float ahrs_A[49] = {0.0f};
static float ahrs_H[21] = {0.0f};
static float ahrs_Q[49] = {0.0f};
static float ahrs_R[9] = {
	0.25f,	0.0f,	0.0f,
	0.0f,	0.25f,	0.0f,
	0.0f,	0.0f,	0.25f
};

static void ahrs_process(void *vfilter);
static void ahrs_measure(void *vfilter, float *estimation);

static void ahrs_computeA(void *vfilter);
static void ahrs_computeQ(void *vfilter);
static void ahrs_acc_computeH(void *vfilter);
static void ahrs_mag_computeH(void *vfilter);
static void ahrs_acc_computeR(void *vfilter);
static void ahrs_mag_computeR(void *vfilter);

struct kalman_filter ahrs_kf = {
	.state_dim = 7,
	.measure_dim = 3,
	.x = ahrs_state,
	.P = ahrs_P,
	.A = ahrs_A,
	.H = ahrs_H,
	.Q = ahrs_Q,
	.R = ahrs_R,
	.process = ahrs_process,
	.measurement = ahrs_measure,
	.computeA = ahrs_computeA,
	.computeQ = ahrs_computeQ,
	.computeH = ahrs_acc_computeH,
	.computeR = ahrs_acc_computeR
};

struct ahrs_monitor {
	uint8_t reset_filter;
	uint64_t last_sensor_time;
	float geomag[3];
};

static struct ahrs_monitor ahrs_monitor = {
	.reset_filter = 0,
	.last_sensor_time = 0
};

void update_ahrs_imu(struct imu_raw_data *imu, struct attitude *att)
{
	float fAcc[3] = { imu->acc[0], imu->acc[1], imu->acc[2] };
	float dt;

	if (att->timestamp == 0) {/* init require */
		acc2angle(fAcc, att->angle);
		angle2quat(att->q, att->angle);
		quat2cbn(att->cbn, att->q);
		att->timestamp = board_get_systime();
	}

	if (ahrs_monitor.reset_filter != 0) {
		ahrs_monitor.reset_filter = 0;
		ahrs_kf.x[0] = att->q[0];
		ahrs_kf.x[1] = att->q[1];
		ahrs_kf.x[2] = att->q[2];
		ahrs_kf.x[3] = att->q[3];
		ahrs_kf.x[4] = 0.0f;
		ahrs_kf.x[5] = 0.0f;
		ahrs_kf.x[6] = 0.0f;

		arm_fill_f32(0.0f, ahrs_kf.P, 49);
		ahrs_kf.P[0] = 1.0f;
		ahrs_kf.P[8] = 1.0f;
		ahrs_kf.P[16] = 1.0f;
		ahrs_kf.P[24] = 1.0f;
		ahrs_kf.P[32] = 0.1f;
		ahrs_kf.P[40] = 0.1f;
		ahrs_kf.P[48] = 0.1f;
	}

	if (ahrs_monitor.last_sensor_time == 0) {
		ahrs_monitor.last_sensor_time = imu->timestamp;
		dt = 0.001f;
	} else {
		dt = (imu->timestamp - ahrs_monitor.last_sensor_time) * 1e-6;
		ahrs_monitor.last_sensor_time = imu->timestamp;
	}

	kf_predict(&ahrs_kf, imu->gyr, dt);

	ahrs_kf.computeH = ahrs_acc_computeH;
	ahrs_kf.computeR = ahrs_acc_computeR;
	kf_update(&ahrs_kf, fAcc);

	att->q[0] = ahrs_kf.x[0];
	att->q[1] = ahrs_kf.x[1];
	att->q[2] = ahrs_kf.x[2];
	att->q[3] = ahrs_kf.x[3];
	quat2angle(att->angle, att->q);
	quat2cbn(att->cbn, att->q);
	att->rate[0] = imu->gyr[0] - ahrs_kf.x[4];
	att->rate[1] = imu->gyr[1] - ahrs_kf.x[5];
	att->rate[2] = imu->gyr[2] - ahrs_kf.x[6];
}

void update_ahrs_mag(struct mag_raw_data *mag, struct attitude *att)
{
	float fMag[3] = {mag->mag[0], mag->mag[1], mag->mag[2]};
	float angle[3] = {att->angle[0], att->angle[1], att->angle[2]};

	mag2angle(fMag, angle);

	if (fabs(wrap_pi(angle[2] - att->angle[2])) > 15.0/180.0*PI) {
		att->angle[2] = angle[2];
		angle2quat(att->q, att->angle);
		ahrs_monitor.reset_filter = 1;
	} else {
		ahrs_kf.computeH = ahrs_mag_computeH;
		ahrs_kf.computeR = ahrs_mag_computeR;

		kf_update(&ahrs_kf, fMag);

		att->q[0] = ahrs_kf.x[0];
		att->q[1] = ahrs_kf.x[1];
		att->q[2] = ahrs_kf.x[2];
		att->q[3] = ahrs_kf.x[3];
		quat2angle(att->angle, att->q);
	}
	quat2cbn(att->cbn, att->q);

	ahrs_monitor.geomag[0] = att->cbn[0]*fMag[0]
		+ att->cbn[1]*fMag[1]
		+ att->cbn[2]*fMag[2];
	ahrs_monitor.geomag[1] = att->cbn[3]*fMag[0]
		+ att->cbn[4]*fMag[1]
		+ att->cbn[5]*fMag[2];
	ahrs_monitor.geomag[2] = att->cbn[6]*fMag[0]
		+ att->cbn[7]*fMag[1]
		+ att->cbn[8]*fMag[2];
	arm_sqrt_f32(ahrs_monitor.geomag[0]*ahrs_monitor.geomag[0]
		     + ahrs_monitor.geomag[1]*ahrs_monitor.geomag[1],
		     &(ahrs_monitor.geomag[0]));
	ahrs_monitor.geomag[1] = 0.0;
}

void ahrs_process(void *vfilter)
{}

void ahrs_measure(void *vfilter, float *estimation)
{}

void ahrs_computeA(void *vfilter)
{}

void ahrs_computeQ(void *vfilter)
{}

void ahrs_acc_computeH(void *vfilter)
{}

void ahrs_mag_computeH(void *vfilter)
{}

void ahrs_acc_computeR(void *vfilter)
{}

void ahrs_mag_computeR(void *vfilter)
{}

